import re
import sys
import os
from collections import defaultdict
from transformers import pipeline, AutoTokenizer, AutoModelForTokenClassification, BertTokenizerFast
import torch
from utils.logger import logger

# 添加项目根目录到路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from config import NER_SENSITIVE_TYPES, NER_MODEL_CACHE_DIR, NER_MODEL_NAME, NER_TOKENIZER_NAME

# 全局缓存模型，避免重复加载
global_cached_pipeline = None

def ensure_model_cache_dir():
    """
    确保模型缓存目录存在
    """
    if not os.path.exists(NER_MODEL_CACHE_DIR):
        logger.info(f"创建模型缓存目录: {NER_MODEL_CACHE_DIR}")
        os.makedirs(NER_MODEL_CACHE_DIR, exist_ok=True)
    return NER_MODEL_CACHE_DIR

def extract_with_model(text):
    """
    使用分块策略处理长文本的NER识别
    """
    global global_cached_pipeline
    
    # 确保缓存目录存在
    cache_dir = ensure_model_cache_dir()
    
    # 初始化模型（仅一次）
    if global_cached_pipeline is None:
        # 检查模型文件是否已存在于缓存目录中
        tokenizer_path = os.path.join(cache_dir, "tokenizer")
        model_path = os.path.join(cache_dir, "model")
        
        if os.path.exists(tokenizer_path) and os.path.exists(model_path):
            # 从本地加载模型
            logger.info(f"从本地缓存加载NER模型: {model_path}")
            try:
                tokenizer = BertTokenizerFast.from_pretrained(tokenizer_path)
                model = AutoModelForTokenClassification.from_pretrained(model_path)
                logger.info(f"成功从本地缓存加载NER模型")
            except Exception as e:
                logger.error(f"从本地加载模型失败: {e}")
                logger.info(f"尝试从在线源重新下载模型...")
                # 如果本地加载失败，尝试在线下载并保存
                tokenizer = BertTokenizerFast.from_pretrained(NER_TOKENIZER_NAME, cache_dir=cache_dir)
                tokenizer.save_pretrained(tokenizer_path)
                model = AutoModelForTokenClassification.from_pretrained(NER_MODEL_NAME, cache_dir=cache_dir)
                model.save_pretrained(model_path)
        else:
            # 首次运行，从在线源下载并保存到本地
            logger.info(f"首次运行，正在下载NER模型到本地缓存: {cache_dir}")
            try:
                # 下载并保存分词器
                tokenizer = BertTokenizerFast.from_pretrained(NER_TOKENIZER_NAME, cache_dir=cache_dir)
                tokenizer.save_pretrained(tokenizer_path)
                logger.info(f"分词器已保存到: {tokenizer_path}")
                
                # 下载并保存模型
                model = AutoModelForTokenClassification.from_pretrained(NER_MODEL_NAME, cache_dir=cache_dir)
                model.save_pretrained(model_path)
                logger.info(f"模型已保存到: {model_path}")
            except Exception as e:
                logger.error(f"下载模型失败: {e}")
                # 如果下载失败，尝试直接在线使用
                logger.info(f"尝试直接在线使用模型...")
                tokenizer = BertTokenizerFast.from_pretrained(NER_TOKENIZER_NAME)
                model = AutoModelForTokenClassification.from_pretrained(NER_MODEL_NAME)
        
        # 创建NER管道
        logger.info(f"正在初始化NER管道...")
        global_cached_pipeline = pipeline(
            "ner", 
            model=model,
            tokenizer=tokenizer,
            aggregation_strategy="simple"
        )
        logger.info(f"NER模型加载完成。")
    
    nlp_pipeline = global_cached_pipeline
    
    # 确定分块参数
    MAX_LEN = 510  # 预留2个位置给特殊token [CLS]和[SEP]
    OVERLAP = 128  # 重叠长度
    
    text_length = len(text)
    
    # 短文本直接处理
    if text_length <= MAX_LEN:
        logger.info(f"文本长度 {text_length} 字符，直接处理")
        ner_results = nlp_pipeline(text)
        return process_entities(ner_results)
    
    # 长文本分块处理
    chunks = []
    start_idx = 0
    
    while start_idx < text_length:
        end_idx = start_idx + MAX_LEN
        if end_idx > text_length:
            end_idx = text_length
        
        # 创建重叠块
        chunk = text[start_idx:end_idx]
        chunks.append((start_idx, chunk))
        
        # 下一步位置考虑重叠
        start_idx += MAX_LEN - OVERLAP
        
        # 防止无限循环
        if start_idx >= text_length:
            break
    
    logger.info(f"文本长度 {text_length} 字符，分为 {len(chunks)} 个块进行处理")
    
    # 处理每个块并收集结果
    all_entities = []
    for i, (start_offset, chunk_text) in enumerate(chunks):
        logger.info(f"处理文本块 {i+1}/{len(chunks)}, 位置: {start_offset}-{start_offset+len(chunk_text)}")
        entities = nlp_pipeline(chunk_text)
        
        # 调整实体位置到原始文本
        for entity in entities:
            entity['start'] += start_offset
            entity['end'] += start_offset
            all_entities.append(entity)
    
    # 实体去重
    deduped_entities = deduplicate_entities(all_entities)
    logger.info(f"NER原始识别实体: {len(all_entities)}个，去重后: {len(deduped_entities)}个")
    return process_entities(deduped_entities)

def process_entities(entities):
    """
    处理实体为统一格式，并根据配置过滤实体类型
    """
    found_entities = defaultdict(list)
    filtered_count = 0
    
    for entity in entities:
        entity_type = entity['entity_group'].upper()
        
        # 根据配置检查该类型是否被标记为敏感信息
        if entity_type in NER_SENSITIVE_TYPES and NER_SENSITIVE_TYPES[entity_type]:
            found_entities[entity_type].append({
                "word": entity['word'],
                "start": entity['start'],
                "end": entity['end'],
                "score": round(entity['score'], 4)
            })
        else:
            filtered_count += 1
    
    # 记录过滤信息
    if filtered_count > 0:
        logger.info(f"根据NER_SENSITIVE_TYPES配置过滤掉 {filtered_count} 个非敏感实体")
    
    # 记录识别结果
    total_entities = sum(len(entities) for entities in found_entities.values())
    logger.info(f"NER最终识别出 {total_entities} 个敏感实体，分布在 {len(found_entities)} 种类型中")
    
    for entity_type, entities_list in found_entities.items():
        logger.info(f"- {entity_type}: {len(entities_list)} 个实体")
    
    return found_entities

def deduplicate_entities(entities):
    """
    去除重叠区域识别的重复实体
    """
    # 按实体类型分组
    grouped = defaultdict(list)
    for e in entities:
        grouped[(e['entity_group'], e['start'], e['end'])].append(e)
    
    # 选择每个位置得分最高的实体
    deduped = []
    for key, group in grouped.items():
        best_entity = max(group, key=lambda x: x['score'])
        deduped.append(best_entity)
    
    return deduped

# -------------------- 步骤三：主函数，整合所有结果 --------------------

def extract_sensitive_info(text):
    """
    使用NER模型提取敏感信息
    返回格式化的实体识别结果
    """
    # 用模型提取
    model_results = extract_with_model(text)
    final_results = defaultdict(list)

    # 添加模型结果
    for entity_type, entities in model_results.items():
        for entity in entities:
            final_results[entity_type].append(entity)

    return final_results


# -------------------- 运行示例 --------------------

if __name__ == "__main__":
    
    # 长文本示例
    # text = ("法狗狗银行股份有限公司固定资产（项目融资）借款合同（2022年第一版）款本息的，就本金部分，乙方有权收取逾期贷款罚息，逾期贷款罚息利率为在本合同约定利率的基础上上浮50％；就利息部分，自结息日次日起对不按时偿付的利息按逾期贷款罚息利率计收复利，即每经过一个结息期后，都要将所生利息加入本金，以计算下期的利息，直至还清为止。前述本金产生的罚息与利息产生的复利分项计算，合计金额为借款人应付款项。本条所称利息均包含罚息。14.3.2.甲方 未按合同用途使用贷款的，乙方有权收取罚息，利率为在本合同约定利率的基础上上浮100％。14.3.3.甲方未按第5.2（分 期的条款）款规定的分期还款计划按期如数还款，对未还贷款部分的贷款利率执行本合同中约定的逾期贷款罚息。14.3.4. 同一贷款存在逾期归还和改变用途两种行为的，按罚息利率中最高罚息利率执行，直至清偿本息止。14.3.5.对于改变用途 和逾期归还外的其他违约行为，从违约之日起，乙方按违约发生日借款余额和违约天数每日向甲方收取万分之五的违约金，直至违约行为纠正之日止；对其中无法纠正的行为，则按违约发生日借款余额一次性收取千分之三的违约金。14.3.6.停止 发放未发放的贷款，对已发放贷款提前收回。14.3.7.解除合同，要求甲方立即归还贷款本息（包括但不限于罚息、复利等 ）。14.3.8.其他依本合同和依法可采取的任何措施。14.4.乙方因追索贷款而发生的费用和损失包括但不限于诉讼费、保全费、仲裁费、律师费(以乙方与律师事务所签订的代理合同约定为准，包括已支付和未支付费用)、公告费、公证费、送达费、评估费、执行费、差旅费、调查费等费用，甲方同意承担。14.5.本合同中所约定的利息、罚息、复利等均为合同结算、 清理条款，即便合同被解除的，亦不影响合同中结算、清理条款的效力。第十五条账户扣划与账户控制权15.1.当发生甲方 违约或乙方提前收回贷款情形时，乙方有权从甲方在法狗狗银行所属各机构开立的账户划收，划收以贷款本金、利息、复利、罚息、违约金及实现债权的费用包括但不限于诉讼费、保全费、仲裁费、律师费(以乙方与律师事务所签订的代理合同约 定为准，包括已支付和未支付费用)、公告费、公证费、送达费、评估费、执行费、差旅费、调查费等为限。该划收视同甲 方已提前书面授权划收，本合同视同划收的不可撤销的书面授权文件。15.2.发生甲方违约或乙方提前收回贷款情形时，乙 方可以对甲方在法狗狗银行所属各机构开立的账户采取账户控制、资金监管等措施，控制、监管范围以本合同项下贷款本金、利息、复利、罚息、违约金及实现债权的费用包括但不限于诉讼费、保全费、仲裁费、律师费(以乙方与律师事务所签订 的代理合同约定为准，包括已支付和未支付费用)、公告费、公证费、送达费、评估费、执行费、差旅费、调查费等为限。 该控制、监管视同甲方已提前书面授权进行，本合同视同控制、监管的不可撤销的书面授权文件。请再次确认上述重要合同条款或合同文本内容，可能影响贵司的重大利益：我方已阅读本合同所有条款，特别是黑体字部分。贵行工作人员已经按照我方的要求就本合同有关我方重大利害关系的条款做了相应的重点说明，并用黑体字进行了充分提示。我方对本合同条款的含义及相应的法律后果已全部知晓并充分理解。除上述内容外，贵行已明确告知我方，不得将信贷资金用于购房等其他非合")
    text=("深圳市龙华区民治街道 白松三路华侨城")
    logger.info(f"文本长度: {len(text)}字符")
    
    # 执行提取任务
    all_found_entities = extract_sensitive_info(text)
    
    # 打印结果
    logger.info("\n识别到的敏感信息如下：")
    for entity_type, entities in all_found_entities.items():
        logger.info(f"--- {entity_type} ({len(entities)}个) ---")
        for entity in entities:  # 最多显示5个
             logger.info(f"  - 文本: {entity['word']}, 位置: ({entity['start']}, {entity['end']}), 置信度: {entity['score']}")
        if len(entities) > 5:
             logger.info(f"  及其他{len(entities)-5}个{entity_type}...")
    
    # 演示脱敏效果
    logger.info("\n脱敏（替换）效果演示：")
    masked_text_list = list(text)
    for entities in all_found_entities.values():
        for entity in entities:
            start, end = entity['start'], entity['end']
            masked_text_list[start:end] = ['*'] * (end - start)
            
    logger.info("".join(masked_text_list))